import numpy as np
from sko.DE import DE
from pyted.frontend import Variable
from pyted import get_circuit, DataSweep, debug

from sko.tools import func_transformer
from torch import Tensor
from frontend.torch import TorchVariable
import sys
import time

class LSHADE():

    def __init__(self, fitness, lowwer, upper, pop_size, dim, epochs, tracer, eval):
        # np.random.seed(1)
        self.fitness = fitness
        # self.obj_func = obj_func
        self.eval = np.array(eval)
        self.tracer = tracer
        self.lowbound = np.array(lowwer)
        self.upbound = np.array(upper)
        self.pop_size = pop_size
        self.dim = dim
        self.population = np.random.uniform(
            self.lowbound, self.upbound, (self.pop_size, self.dim))
        self.fit = np.random.rand(self.pop_size)
        self.conv = np.random.rand(self.pop_size)
        self.best = self.population[0]
        # self.mut_way = mut_way
        self.Archive = np.empty((0, self.dim))
        self.time_up = 2*60*60 - 200
        self.start_time = time.time()

        self.uCR = 0.5
        self.uF = 0.5
        self.F = np.array([0.5 for _ in range(self.pop_size)])
        self.CR = np.array([0.5 for _ in range(self.pop_size)])
        self.c = 0.01
        self.SCR = []
        self.SF = []
        self.success_set = []
        self.fail_set = []
        self.H = 6
        self.MCR = np.array([0.5] * self.H)
        self.MF = np.array([0.5] * self.H)
        self.stopvalue = 0
        # ***************************************
        self.Epochs = epochs
        self.Ninit = self.pop_size
        self.Nmin = 4
        self.NFE = 0
        self.max_NFE = 10000*self.dim

    def initpop(self):
        self.population = self.lowbound + \
            (self.upbound - self.lowbound) * \
            np.random.rand(self.pop_size-1, self.dim)
        self.population = np.concatenate((self.population, self.eval.reshape(1,-1)), axis=0)
        #with SuppressOutput():
        self.fit, self.conv = self.fitness(self.population)
        self.NFE += self.pop_size
        # self.best = self.population[np.argmin(self.fit)]

    def mut(self):
        p = max(int(self.pop_size * 0.1), 1)
        # idx_set = np.argsort(self.fit)[:p]
        idx_set = np.lexsort((self.conv, self.fit))[:p]
        xpbest_idx = np.random.choice(idx_set, self.pop_size, replace=True)
        xpbest = self.population[xpbest_idx].reshape(
            -1, self.population.shape[1])
        a_indices = np.random.choice(
            np.delete(np.arange(self.pop_size), xpbest_idx), self.pop_size, replace=True)
        a = self.population[a_indices].reshape(-1, self.population.shape[1])
        combined_population = np.concatenate(
            (self.population, self.Archive), axis=0)
        b_indices = np.random.choice(
            len(combined_population), self.pop_size, replace=True)
        b = combined_population[b_indices]
        F_broadcasted = np.tile(self.F[:, np.newaxis], (1, self.dim))
        v = self.population + F_broadcasted * \
            (xpbest - self.population) + F_broadcasted * (a - b)

        return v

    def cross(self, mut_chorm):
        j = np.random.randint(0, self.dim, size=self.pop_size)
        CR_broadcasted = np.tile(self.CR[:, np.newaxis], (1, self.dim))
        crossover_mask = (np.random.rand(self.pop_size, self.dim) < CR_broadcasted) | (
            np.arange(self.dim) == j[:, np.newaxis])
        cross_chorm = np.where(crossover_mask, mut_chorm, self.population)
        # cross_chorm = np.clip(cross_chorm, self.lowbound, self.upbound)

        nan_mask = np.isnan(cross_chorm)
        out_of_bounds = (cross_chorm < self.lowbound) | (
            cross_chorm > self.upbound) | nan_mask
        new_values = np.random.uniform(
            self.lowbound, self.upbound, size=cross_chorm.shape)
        cross_chorm[out_of_bounds] = new_values[out_of_bounds]

        return cross_chorm

    def select(self, cross_chorm):
        #with SuppressOutput():
        temp, temp_v = self.fitness(cross_chorm)
        self.NFE += self.pop_size
        mask_1 = (self.conv == temp_v) & (self.fit >= temp)
        mask_2 = (self.conv > temp_v)
        mask = mask_1 | mask_2

        self.SCR.extend(self.CR[mask])
        self.SF.extend(self.F[mask])
        self.success_set.extend(np.where(mask)[0])
        self.fail_set.extend(self.fit[mask])
        self.Archive = np.concatenate(
            (self.Archive, self.population[mask]), axis=0)

        if self.Archive.shape[0] > self.pop_size * 2.6:
            num_to_delete = int(self.Archive.shape[0] - self.pop_size * 2.6)
            indices_to_delete = np.random.choice(
                self.Archive.shape[0], num_to_delete, replace=False)

            self.Archive = np.delete(self.Archive, indices_to_delete, axis=0)

        # update_mask = (temp <= self.fit) & mask
        self.population[mask] = cross_chorm[mask]
        self.fit[mask] = temp[mask]
        self.conv[mask] = temp_v[mask]

    def Linear_reduction(self, epoch):
        resize = int(self.Ninit-self.NFE/(20000*self.dim)
                     * (self.Ninit-self.Nmin))
        if resize < 20:
            return
        self.pop_size = resize
        idx = np.argsort(self.fit)[:self.pop_size]
        self.population = self.population[idx]
        # self.pop_size = self.population.shape[0]
        self.fit = self.fit[idx]
        self.conv = self.conv[idx]

    def run(self):
        self.initpop()
        min_constraint_indices = np.where(self.conv == np.min(self.conv))[0]
        min_fitness_among_min_constraint = np.argmin(self.fit[min_constraint_indices])
        min_fitness_index = min_constraint_indices[min_fitness_among_min_constraint]
        self.best = self.population[min_fitness_index]
        min_fitness_value = self.fit[min_fitness_index]
        min_constraint_value = self.conv[min_fitness_index]
        self.tracer.collect_n_iters_fom(
                min_fitness_value,
                np.array([min_constraint_value]),
            )
        if min_constraint_value == 0:
            return self.best, self.tracer
        for epoch in range(self.Epochs):
            self.SCR = []
            self.SF = []
            self.success_set = []
            self.fail_set = []
            mcr_indices = np.random.choice(len(self.MCR), self.pop_size)
            mf_indices = np.random.choice(len(self.MF), self.pop_size)
            self.CR = np.where(self.MCR[mcr_indices] == self.stopvalue, 0, np.random.normal(
                self.MCR[mcr_indices], 0.1))
            self.CR = np.clip(self.CR, 0, 1)
            self.F = self.MF[mf_indices] + \
                np.sqrt(0.1) * np.random.standard_cauchy(self.pop_size)
            self.F[self.F <= 0] = np.random.standard_cauchy(
                np.sum(self.F <= 0))
            self.F[self.F > 1] = 1
            mut_chorms = self.mut()
            cross_chorms = self.cross(mut_chorms)
            self.select(cross_chorms)

            # self.best = self.population[np.argmin(self.fit)]
            # min_conv = np.min(self.conv)
            # min_fit = np.min(self.fit[self.conv == min_conv])
            min_constraint_indices = np.where(
                self.conv == np.min(self.conv))[0]
            min_fitness_among_min_constraint = np.argmin(
                self.fit[min_constraint_indices])
            min_fitness_index = min_constraint_indices[min_fitness_among_min_constraint]
            self.best = self.population[min_fitness_index]
            min_fitness_value = self.fit[min_fitness_index]
            min_constraint_value = self.conv[min_fitness_index]
            self.tracer.collect_n_iters_fom(
                    min_fitness_value,
                    np.array([min_constraint_value]),
                )
            if min_constraint_value == 0:
                # print(min_constraint_value, min_fitness_value)
                break
            if (np.any(self.SCR) and np.any(self.SF)):
                d_f = np.abs(self.fail_set - self.fit[self.success_set])
                w = d_f / np.sum(d_f)
                self.uCR = np.sum(np.multiply(w, np.array(
                    self.SCR) ** 2)) / np.sum(np.multiply(w, np.array(self.SCR)))
                self.uF = np.sum(np.multiply(w, np.array(self.SF) ** 2)) / \
                    np.sum(np.multiply(w, np.array(self.SF)))
                self.MCR[epoch % self.H] = self.stopvalue if (
                    self.MF[epoch % self.H] == self.stopvalue or max(self.SCR) == 0) else self.uCR
                self.MF[epoch % self.H] = self.uF
            self.Linear_reduction(epoch)
            # ft, cv = self.fitness(self.best)
            # print(np.min(self.conv), end='\t')
            # print('epoch:', epoch, 'best:', np.min(self.fit),
            #       'mean:', np.mean(self.fit),
            #       'worst:', np.max(self.fit),
            #       'std:', np.std(self.fit),
            #       #   'NFE:', self.NFE,
            #       'pop_size', self.pop_size,
            #       "min_fit:", min_fitness_value,
            #       "min_conv:", min_constraint_value,
            #       'time:', time.time()-self.start_time)
            if time.time() - self.start_time > self.time_up:
                break
        # print(self.best)
        # print(type(self.best))
        # print(self.best.shape)
        # tmp = self.obj_func(self.best.reshape(1,6))
        # tmp = self.constraints(self.best.reshape(1,6))
        return self.best, self.tracer  # ·µ»Ø×îÓÅ¸öÌå


class My_Optimizer:
    def __init__(
        self,
        tester,
        problem,
        tracer=None,
        ref_point=None,
        pop_size=50,
        maxiters=1000,
    ):
        # set the parameters
        # =============================#
        # =============================#
        self.problem = problem
        self.tester = tester
        self.pop_size = pop_size
        self.maxiters = maxiters
        self.tracer = tracer(ref_point)
        Dim = len(get_circuit().variables)
        bound = [(v._lb, v._ub, v._eval) for _, v in get_circuit().variables.items()]
        lb, ub, eval = map(list, zip(*bound))

        def obj_func(x):
            # # WRITE CODE HERE within two '#' bar                                        #
            # # Converts the input x to list form and dim(x)=1                            #
            # #############################################################################

            # #############################################################################
            tester = self.update(x)
            prob = self.problem(tester)

            self._obj = prob.objective[0]
            obj = self._obj.clone().detach().numpy()
            return obj

        def obj_func_parallel(params):
            # # WRITE CODE HERE within two '#' bar 						#
            # # Converts the input params to list x form and dim(x)=2 	#
            # ###########################################################
            x = params
            # ###########################################################
            keylib = []
            reslib = []
            newlib = []
            f = []
            CV = []
            params = list(get_circuit().variables.keys())
            get_circuit().datasweep = DataSweep("sweepdata", params, x.T)
            with get_circuit().setenv(shell=False):
                res = self.tester.build()
                for k, v in res.items():
                    keylib.append(k)
                    reslib.append(v)
                reslib = Tensor(reslib).T
                for j in range(len(reslib)):
                    newdict = {}
                    for index, value in enumerate(keylib):
                        newdict[value] = TorchVariable(reslib[j][index])
                    newlib.append(newdict)
            test = [self.problem(z) for z in newlib]
            f_ret = [x.objective[0] for x in test]
            cv = [[list(x.constraint)] for x in test]
            CV = np.vstack(cv)
            f = np.array(f_ret)
            penaly_ueq = np.maximum(CV, 0).sum(axis=1)
            return f, penaly_ueq

        def constrains(x):
            CV = []
            for i in range(len(x)):
                tester = self.update(x[i])
                prob = self.problem(tester)
                cv = [x.tolist() for x in prob.constraint]
                CV.append(cv)
                if self.tracer != None:
                    self.tracer.collect_n_iters_fom(
                        prob.objective[0].clone().detach().numpy(),
                        np.array([con.to_numpy() for con in prob.constraint]),
                    )

            CV = np.vstack(CV)
            penaly_ueq = np.maximum(CV, 0).sum(axis=1)
            return penaly_ueq

        self.method = LSHADE(
            fitness=obj_func_parallel,
            lowwer=lb,
            upper=ub,
            pop_size=50,
            dim=Dim,
            epochs=1000000,
            tracer=self.tracer,
            eval=eval
        )

    def solve(self, prophetVars=None):
        '''
        prophetVars:
        '''
        best_x, tracer_t = self.method.run()
        iters_info = (
            (tracer_t.n_iters, tracer_t.n_iters_fom)
            if self.tracer != None
            else (None, None)
        )

        return {
            "success": True,
            "Vars": [best_x],
            "fom": iters_info,
        }

        pass

    def update(self, x):
        Variable.update_all_variables(x)
        with get_circuit().setenv(quite=True, shell=False):
            return self.tester.build()